#define FLEXIO_DEV_VER_USED FLEXIO_DEV_VER(24, 9, 0)

#include <doca_dpa_dev.h>
#include <doca_dpa_dev_buf.h>
#include "wrapper_flexio_device.h"
#include "common_cross.h"
#include "feature_cross.h"

static struct DpaThreadInfo dpa_execute_unit[100010];
/*
 * Kernel function for kernel_launch sample
 */
void test_func(void) {
	DOCA_DPA_DEV_LOG_INFO ("are you ok?");
}

__dpa_global__ void hello_world(doca_dpa_dev_mmap_t mmap_handle, uint64_t addr)
{
	DOCA_DPA_DEV_LOG_INFO ("%d\n", doca_dpa_dev_thread_rank());
	if (doca_dpa_dev_thread_rank() == 0) {
		DOCA_DPA_DEV_LOG_INFO("Hello from kernel\n");
		doca_dpa_dev_uintptr_t dpa_addr = doca_dpa_dev_mmap_get_external_ptr(mmap_handle, addr);
		int64_t * dpa_ptr = (int64_t *) dpa_addr;
		DOCA_DPA_DEV_LOG_INFO("read from dpa value is %ld\n", (*dpa_ptr));
		(*dpa_ptr) = 111;
		test_func();
	}
}

void dpa_fillnull(uint64_t *input_data, uint64_t *output_data, struct DpaThreadInfo *arg) {
    int64_t size = arg->size;

    for (int i = 0; i < size; ++i) {
        int64_t x = ((int64_t*)input_data)[i];
        if (x == 1) x = 2;
        ((int64_t*)output_data)[i] = x;
    }
}

void dpa_cast(uint64_t *input_data, uint64_t *output_data, struct DpaThreadInfo *arg) {
    int64_t size = arg->size;

    for (int i = 0; i < size; ++i) {
        int64_t x = ((int64_t*)input_data)[i];
        ((int64_t*)output_data)[i] = x;
    }
}

void dpa_sigrid_hash(uint64_t *input_data, uint64_t *output_data, struct DpaThreadInfo *arg) {
    int64_t size = arg->size;
    int salt = SIGRID_ARG1;
    int num_embeddings = SIGRID_ARG2;

    for (int i = 0; i < size; ++i) {
        int64_t x = ((int64_t*)input_data)[i];
        uint64_t value = (x ^ (uint64_t)(salt));
        value = (value * 2654435761U) ^ (value >> 16);
        value = value % (uint64_t)(num_embeddings);
        ((int64_t*)output_data)[i] = value;
    }
}

void dpa_clamp_list(uint64_t *input_data, uint64_t *output_data, struct DpaThreadInfo *arg) {
    int64_t size = arg->size;

    for (int i = 0; i < size; ++i) {
        int64_t x = ((int64_t*)input_data)[i];
        if (x < CLAMP_ARG1) x = CLAMP_ARG1;
        else if (x > CLAMP_ARG2) x = CLAMP_ARG2;
        ((int64_t*)output_data)[i] = x;
    }
}

void dpa_ngram_list(uint64_t *input_data, uint64_t *output_data, struct DpaThreadInfo *arg) {
    int64_t size = arg->size;
    int num_ele_per_feature = arg->num_ele_per_feature;
    int32_t actual_n_of_gram = NGRAM_ARG > num_ele_per_feature ? num_ele_per_feature : NGRAM_ARG;
    int32_t final_length = (num_ele_per_feature - actual_n_of_gram + 1) * actual_n_of_gram;

    for (int i = 0; i < size; ++i) {
        int cnt_output = 0;
        for (int k = 0; k < num_ele_per_feature - actual_n_of_gram + 1; k++) {
            for (int l = 0; l < actual_n_of_gram; l++) {
                int input_index = i * num_ele_per_feature + k + l;
                int64_t x = ((int64_t*)input_data)[input_index];

                int output_index = i * final_length + cnt_output;
                ((int64_t*)output_data)[output_index] = x;

                ++cnt_output;
            }
        }
    }
}

void dpa_firstx(uint64_t *input_data, uint64_t *output_data, struct DpaThreadInfo *arg) {
    int64_t size = arg->size;
    int num_ele_per_feature = arg->num_ele_per_feature;

    for (int i = 0; i < size; ++i) {
        for (int k = 0; k < (int)FIRSTX_ARG; ++k) {
            int input_index = i * num_ele_per_feature + k;
            int64_t x = ((int64_t*)input_data)[input_index];
            int output_index = i * FIRSTX_ARG + k;
            ((int64_t*)output_data)[output_index] = x;
        }
    }
}

void dpa_logit(int stream_id, int *data, int size) {
	(void)stream_id;
    for (int i = 0; i < size / (int)sizeof(int); ++i) {
        data[i] = data[i] * 2;
		// test
		data[i] = data[i] % 10;
    }
}

__attribute__((unused))
void print_node(struct DpaThreadInfo * node) {
	printf ("size : %d\n, num_ele : %d, input_lkey : %d, output_lkey : %d\n",
		node->size, node->num_ele_per_feature, node->input_mem_mkey, node->output_mem_mkey);
	printf ("mem_input_addr : %ld, mem_output_addr : %ld\n", (int64_t)node->input_mem_addr, (int64_t)node->output_mem_addr);
	printf ("windows_id = %d\n", node->window_id);
	printf ("preprocess_type = %d\n", node->preprocess_type);
}

__dpa_rpc__ uint64_t dpa_execute_unit_init(uint64_t data) {   
	struct DpaThreadInfo *unit_data    = (struct DpaThreadInfo *)data;
	struct DpaThreadInfo Tmp = (*unit_data);

	int batch_dim  = Tmp.size;
	int thread_num = Tmp.thread_num;
	printf ("current init thread_num = %d\n", thread_num);
	// int thread_index = Tmp.thread_index;
	int step = batch_dim / thread_num;
	int remainder = batch_dim % thread_num;
	int num_ele_per_feature = Tmp.num_ele_per_feature;
	int output_num_ele_per_feature = Tmp.output_num_ele_per_feature;

	for (int j = 0; j < thread_num; ++j) {
		struct DpaThreadInfo T;
		T = Tmp;

		int res = j;
		if (res > remainder) res = remainder;
		int start_offset = j * step + res;
		int end_offset = start_offset + step + (j < remainder ? 1 : 0);
		int num_elements = end_offset - start_offset;

		T.input_mem_addr  = (char*) T.input_mem_addr + start_offset * num_ele_per_feature * sizeof(int64_t);
        T.output_mem_addr = (char*) T.output_mem_addr + start_offset * output_num_ele_per_feature * sizeof(int64_t);

		T.input_mem_offset  = start_offset * num_ele_per_feature * sizeof(int64_t);
		T.output_mem_offset = start_offset * output_num_ele_per_feature * sizeof(int64_t);
		
		T.thread_index += j;
		T.size = num_elements;
		T.stream_id = j;
		dpa_execute_unit[T.thread_index] = T;
	}

	__dpa_thread_system_fence();

	return 0;
}

size_t time_collect[256];

__dpa_global__ void docasdkprocess(doca_dpa_dev_mmap_t mmap_handle, uint64_t arg_daddr) {
	// return;
	size_t start = __dpa_thread_cycles();
	struct DpaThreadInfo *arg = (struct DpaThreadInfo *) &dpa_execute_unit[arg_daddr + doca_dpa_dev_thread_rank()];
	
	// printf ("recieve arg_daddr = %lu\n", arg_daddr);
	// printf ("thread_num = %d\n", arg->thread_num);
	// printf ("num_ele_per_feature = %d\n", arg->num_ele_per_feature);

	uint64_t *input_data = (uint64_t *) doca_dpa_dev_mmap_get_external_ptr(mmap_handle, (uint64_t)arg->input_mem_addr);
	uint64_t *output_data = (uint64_t *) doca_dpa_dev_mmap_get_external_ptr(mmap_handle, (uint64_t)arg->output_mem_addr);

	switch (arg->preprocess_type) {
		case LOGIT:
			dpa_logit(1, (int *) input_data, 1); // deprecate
			// TODO: 实现 LOGIT 操作
			break;
		case BUCKETIZE:
			// TODO: 实现 BUCKETIZE 操作
			break;
		case SIGRID_HASH:
			dpa_sigrid_hash(input_data, output_data, arg);
			// TODO: 实现 SIGRID_HASH 操作
			break;
		case CLAMP_LIST:
			// TODO: 实现 CLAMP_LIST 操作
			dpa_clamp_list(input_data, output_data, arg);
			break;
		case BOXCOX:
			// TODO: 实现 BOXCOX 操作
			break;
		case ONEHOT:
			// TODO: 实现 ONEHOT 操作
			break;
		case NGRAM:
			// TODO: 实现 NGRAM 操作
			dpa_ngram_list(input_data, output_data, arg);
			break;
		case FIRSTX:
			// TODO: 实现 FIRSTX 操作
			dpa_firstx(input_data, output_data, arg);
			break;
		case MAPID:
			// TODO: 实现 MAPID 操作
			break;
		case FILLNULL:
			dpa_fillnull(input_data, output_data, arg);
			// TODO: 实现 FILLNULL 操作
			break;
		case CAST:
			// TODO: 实现 CAST 操作
			break;
		default:
			// TODO: 未知操作类型，可能报错或忽略
			break;
	}
	__dpa_thread_window_writeback();
	size_t end = __dpa_thread_cycles();
	if (doca_dpa_dev_thread_rank() == 0) printf ("total_cycles : %lu\n", end - start);
	fflush(stdout);
}

static uint64_t external_addr;

// int check_thread_num[256];
size_t external_read_time[256];
size_t dpa_mem_read_time[256];
size_t operation_time[256];
size_t atomic_store_time[256];

__dpa_rpc__ uint64_t set_addr(uint64_t daddr) {
	external_addr = daddr;
	return 0;
}

__dpa_rpc__ uint64_t sequence_preprocess(uint64_t task_id) {

	size_t start = __dpa_thread_cycles();
	struct external_ptr * cur_ptr_info = (struct external_ptr *) external_addr;
	void * output_mem_addr_test = cur_ptr_info->output_mem_addr;
	void * input_mem_addr_test  = cur_ptr_info->input_mem_addr;

	// struct flexio_dev_thread_ctx *dtctx;
	// Assert(flexio_dev_get_thread_ctx(&dtctx) >= 0);

	// uint32_t thread_id = flexio_dev_get_thread_id(dtctx);
	struct DpaThreadInfo *arg = (struct DpaThreadInfo *) &dpa_execute_unit[task_id];
	// print_node(arg);
	// int stream_id = arg->stream_id;
	// LOG_F(stream_id, "Bench func, with stream id:%d\n", stream_id);

	uint64_t *input_data  = NULL;
	uint64_t *output_data = NULL;

	size_t start_in = __dpa_thread_cycles();

	output_data = (uint64_t *)get_host_buffer(arg->window_id, arg->output_mem_mkey, (void *)( (char*)output_mem_addr_test + arg->output_mem_offset) );
	input_data  = (uint64_t *)get_host_buffer(arg->window_id, arg->input_mem_mkey, (void *) ( (char*)input_mem_addr_test + arg->input_mem_offset) );

	size_t end_in = __dpa_thread_cycles();
	dpa_mem_read_time[arg->stream_id] = end_in - start_in;

	// output_data = (uint64_t *)get_host_buffer(arg->window_id, arg->output_mem_mkey, arg->output_mem_addr);
	// input_data  = (uint64_t *)get_host_buffer(arg->window_id, arg->input_mem_mkey, arg->input_mem_addr);

	// size_t end_data = __dpa_thread_cycles();
	// printf ("read total_cycles : %lu\n", end_data - start);
	// fflush(stdout);
	// LOG_F(arg->stream_id, "read thread_cycle = %lu\n", end_data - start);

	size_t start_op = __dpa_thread_cycles();
	switch (arg->preprocess_type) {
		case LOGIT:
			dpa_logit(1, (int *) input_data, 1); // deprecate
			// TODO: 实现 LOGIT 操作
			break;
		case BUCKETIZE:
			// TODO: 实现 BUCKETIZE 操作
			break;
		case SIGRID_HASH:
			dpa_sigrid_hash(input_data, output_data, arg);
			// TODO: 实现 SIGRID_HASH 操作
			break;
		case CLAMP_LIST:
			// TODO: 实现 CLAMP_LIST 操作
			dpa_clamp_list(input_data, output_data, arg);
			break;
		case BOXCOX:
			// TODO: 实现 BOXCOX 操作
			break;
		case ONEHOT:
			// TODO: 实现 ONEHOT 操作
			break;
		case NGRAM:
			// TODO: 实现 NGRAM 操作
			dpa_ngram_list(input_data, output_data, arg);
			break;
		case FIRSTX:
			// TODO: 实现 FIRSTX 操作
			dpa_firstx(input_data, output_data, arg);
			break;
		case MAPID:
			// TODO: 实现 MAPID 操作
			break;
		case FILLNULL:
			dpa_fillnull(input_data, output_data, arg);
			// TODO: 实现 FILLNULL 操作
			break;
		case CAST:
			// TODO: 实现 CAST 操作
			dpa_cast(input_data, output_data, arg);
			break;
		default:
			// TODO: 未知操作类型，可能报错或忽略
			break;
	}

	size_t end_op = __dpa_thread_cycles();
	operation_time[arg->stream_id] = end_op - start_op;

	// // if (arg->preprocess_type == LOGIT) {
    // //     dpa_logit(input_data, output_data, arg);
    // // }
	// size_t end = __dpa_thread_cycles();
	// if (arg->stream_id == 0) {
	// 	printf ("total_cycles : %lu\n", end - start);
	// 	fflush(stdout);
	// }

	// output_data[0] = (uint64_t)thread_id;
	// LOG_F(arg->stream_id, "thread_cycle = %lu\n", end - start);

	__dpa_thread_window_writeback();
	size_t end = __dpa_thread_cycles();
    time_collect[arg->stream_id] = end - start;
	return 0;
}

#include <stdatomic.h>

// static atomic_int threads_completed = 0;

static struct ThreadControl threadcontrol; 

static atomic_int public_key = 0;

void dpa_test(uint64_t *input_data, uint64_t *output_data, struct DpaThreadInfo *arg) {
    int64_t size = arg->size;
	(void) input_data;

    for (int i = 0; i < size; ++i) {
        // int64_t x = ((int64_t*)input_data)[i];
		// public_key ++;
        ((int64_t*)output_data)[i] = 8;
    }
}

__dpa_rpc__ uint64_t threadcontrol_init(uint64_t data) {
	struct ThreadControl *unit_data    = (struct ThreadControl *)data;
	threadcontrol = (*unit_data);
	return 0;
}

__dpa_rpc__ uint64_t test_preprocess(uint64_t arg_daddr) {
	// size_t start = __dpa_thread_cycles();

	// struct flexio_dev_thread_ctx *dtctx;
	// Assert(flexio_dev_get_thread_ctx(&dtctx) >= 0);

	// uint32_t thread_id = flexio_dev_get_thread_id(dtctx);

	uint64_t * thread_data = NULL;
	thread_data = (uint64_t *)get_host_buffer(threadcontrol.window_id, threadcontrol.mem_mkey, threadcontrol.mem_addr);
	uint64_t task_id = 0;
	while (1) {
		__dpa_thread_window_read_inv();
		while (1) {
			if ( ( (volatile uint64_t*) thread_data)[0] != task_id){
				break;
			}

			__dpa_thread_window_read_inv();
		}

		size_t start = __dpa_thread_cycles();

		size_t start_ex = __dpa_thread_cycles();
		
		task_id = ((volatile uint64_t*)thread_data)[0];
		
		struct DpaThreadInfo *arg = (struct DpaThreadInfo *) &dpa_execute_unit[task_id + arg_daddr];

		size_t end_ex = __dpa_thread_cycles();
		external_read_time[arg_daddr] = end_ex - start_ex;

		size_t start_in = __dpa_thread_cycles();

		uint64_t *input_data  = NULL;
		uint64_t *output_data = NULL;

		output_data = (uint64_t *)get_host_buffer(arg->window_id, arg->output_mem_mkey, arg->output_mem_addr);
		input_data  = (uint64_t *)get_host_buffer(arg->window_id, arg->input_mem_mkey, arg->input_mem_addr);

		size_t end_in = __dpa_thread_cycles();
		dpa_mem_read_time[arg_daddr] = end_in - start_in;

		size_t start_op = __dpa_thread_cycles();

		switch (arg->preprocess_type) {
			case LOGIT:
				dpa_logit(1, (int *) input_data, 1); // deprecate
				// TODO: 实现 LOGIT 操作
				break;
			case BUCKETIZE:
				// TODO: 实现 BUCKETIZE 操作
				break;
			case SIGRID_HASH:
				dpa_sigrid_hash(input_data, output_data, arg);
				// TODO: 实现 SIGRID_HASH 操作
				break;
			case CLAMP_LIST:
				// TODO: 实现 CLAMP_LIST 操作
				dpa_clamp_list(input_data, output_data, arg);
				break;
			case BOXCOX:
				// TODO: 实现 BOXCOX 操作
				break;
			case ONEHOT:
				// TODO: 实现 ONEHOT 操作
				break;
			case NGRAM:
				// TODO: 实现 NGRAM 操作
				dpa_ngram_list(input_data, output_data, arg);
				break;
			case FIRSTX:
				// TODO: 实现 FIRSTX 操作
				dpa_firstx(input_data, output_data, arg);
				break;
			case MAPID:
				// TODO: 实现 MAPID 操作
				dpa_test(input_data, output_data, arg);
				break;
			case FILLNULL:
				dpa_fillnull(input_data, output_data, arg);
				// TODO: 实现 FILLNULL 操作
				break;
			case CAST:
				// TODO: 实现 CAST 操作
				dpa_cast(input_data, output_data, arg);
				break;
			default:
				// TODO: 未知操作类型，可能报错或忽略
				break;
		}

		size_t end_op = __dpa_thread_cycles();
		operation_time[arg_daddr] = end_op - start_op;

		size_t start_at = __dpa_thread_cycles();
		// if (arg->stream_id == 0) printf ("total_cycles : %lu\n", end - start);
		// _Atomic(uint64_t)* atomic_ptr = (_Atomic(uint64_t)*)output_data;
		
		// 执行原子加法
		// atomic_fetch_add(atomic_ptr, 1);

		// atomic_fetch_add(&threads_completed, 1);
		// int safe_value = atomic_load(&threads_completed);
		// output_data[0] = (uint64_t)safe_value;
		// output_data[0] = (uint64_t)1;

		

		// size_t end = __dpa_thread_cycles();
		// if (arg->stream_id == 0) printf ("total_cycles : %lu\n", end - start);
		__dpa_thread_window_writeback();

		__dpa_thread_system_fence();

		int old_value = atomic_fetch_add(&public_key, 1);

		__dpa_thread_system_fence();

		if (old_value == arg->thread_num - 1){
			atomic_store(&public_key, 0);
			__dpa_thread_system_fence();
			((volatile uint64_t*)thread_data)[1] = arg->thread_num;
		}

		__dpa_thread_window_writeback();
		
		size_t end_at = __dpa_thread_cycles();
		atomic_store_time[arg_daddr] = end_at - start_at;
		
		size_t end = __dpa_thread_cycles();
		time_collect[arg_daddr] = end - start;
		__dpa_thread_system_fence();
		// break;
	}

	if (arg_daddr == 0) {

	}

	return 0;
}

__dpa_rpc__ uint64_t get_max_thread_time(void) {
    size_t max_cycles = 0, pos = 0;
    for (int i = 0; i < 256; ++i) {
        if (time_collect[i] > max_cycles) {
			max_cycles = time_collect[i];
			pos = i;
		}
		// printf ("cycle[%d] = %lu\n", i, time_collect[i]);
		// fflush(stdout);
    }
	printf ("normal_mem_access cycle = %lu\n", external_read_time[pos]);
	printf ("get_host_mem_addr cycle = %lu\n", dpa_mem_read_time[pos]);
	printf ("operation_time cycle = %lu\n", operation_time[pos]);
	printf ("atomic_store_time cycle = %lu\n", atomic_store_time[pos]);
	printf ("total_time cycle = %lu\n", time_collect[pos]);
	fflush(stdout);

    return max_cycles;
}
